// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/benchmark.h>
#include <melon/init/init.h>

#include <turbo/flags/flag.h>

#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/vector/decoded_vector.h>
#include <pollux/vector/vector_fuzzer.h>

TURBO_FLAG(int64_t,fuzzer_seed, 99887766, "Seed for random input dataset generator");

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::test;

namespace {

class SelectivityVectorBenchmark
    : public functions::test::FunctionBenchmarkBase {
 public:
  explicit SelectivityVectorBenchmark(size_t vectorSize)
      : FunctionBenchmarkBase(),
        vectorSize_(vectorSize),
        rowsAll_(vectorSize),
        rows99PerCent_(vectorSize),
        rows50PerCent_(vectorSize),
        rows10PerCent_(vectorSize),
        rows1PerCent_(vectorSize) {
    VectorFuzzer::Options opts;
    opts.vectorSize = vectorSize_;
    opts.nullRatio = 0;
    VectorFuzzer fuzzer(opts, pool(), FLAGS_fuzzer_seed);
    flatVector_ = fuzzer.fuzzFlat(BIGINT());

    for (size_t i = 0; i < vectorSize_; ++i) {
      // Set half to invalid.
      if (fuzzer.coinToss(0.5)) {
        rows50PerCent_.setValid(i, false);
      }

      // Set 90% to invalid.
      if (fuzzer.coinToss(0.9)) {
        rows10PerCent_.setValid(i, false);
      }

      // Set 99% to invalid.
      if (fuzzer.coinToss(0.99)) {
        rows1PerCent_.setValid(i, false);
      }

      // Set 1% to invalid.
      if (fuzzer.coinToss(0.01)) {
        rows99PerCent_.setValid(i, false);
      }
    }

    rowsAll_.updateBounds();
    rows99PerCent_.updateBounds();
    rows50PerCent_.updateBounds();
    rows10PerCent_.updateBounds();
    rows1PerCent_.updateBounds();
  }

  size_t runBaseline() {
    const int64_t* flatBuffer = flatVector_->values()->as<int64_t>();
    size_t sum = 0;

    for (auto i = 0; i < vectorSize_; i++) {
      sum += flatBuffer[i];
    }
    melon::doNotOptimizeAway(sum);
    return vectorSize_;
  }

  size_t runSelectivityAll() {
    return run(rowsAll_);
  }

  size_t runSelectivity50PerCent() {
    return run(rows50PerCent_);
  }

  size_t runSelectivity10PerCent() {
    return run(rows10PerCent_);
  }

  size_t runSelectivity1PerCent() {
    return run(rows1PerCent_);
  }

  size_t runSelectivity99PerCent() {
    return run(rows99PerCent_);
  }

 private:
  size_t run(const SelectivityVector& rows) {
    const int64_t* flatBuffer = flatVector_->values()->as<int64_t>();
    size_t sum = 0;
    rows.applyToSelected(
        [&flatBuffer, &sum](auto row) { sum += flatBuffer[row]; });
    melon::doNotOptimizeAway(sum);
    return vectorSize_;
  }

  const size_t vectorSize_;
  VectorPtr flatVector_;

  SelectivityVector rowsAll_;
  SelectivityVector rows99PerCent_;
  SelectivityVector rows50PerCent_;
  SelectivityVector rows10PerCent_;
  SelectivityVector rows1PerCent_;
};

std::unique_ptr<SelectivityVectorBenchmark> benchmark;

template <typename Func>
void run(Func&& func, size_t iterations = 100) {
  for (auto i = 0; i < iterations; i++) {
    func();
  }
}

BENCHMARK(sumBaselineAll) {
  run([] { benchmark->runBaseline(); });
}

BENCHMARK(sumSelectivityAll) {
  run([] { benchmark->runSelectivityAll(); });
}

BENCHMARK(sumSelectivity99PerCent) {
  run([] { benchmark->runSelectivity99PerCent(); });
}

BENCHMARK(sumSelectivity50PerCent) {
  run([] { benchmark->runSelectivity50PerCent(); });
}

BENCHMARK(sumSelectivity10PerCent) {
  run([] { benchmark->runSelectivity10PerCent(); });
}

BENCHMARK(sumSelectivity1PerCent) {
  run([] { benchmark->runSelectivity1PerCent(); });
}

} // namespace

int main(int argc, char* argv[]) {
  melon::Init init{&argc, &argv};
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  memory::MemoryManager::initialize({});
  benchmark = std::make_unique<SelectivityVectorBenchmark>(10'000);
  melon::runBenchmarks();
  benchmark.reset();
  return 0;
}
